from PIL import Image
import numpy as np
import os
import cv2
import copy
import io
import random

def split_image(img_path, diff_height=0):
    txt_img = cv2.imread(img_path)

    # print(txt_img)
    # cv2.imwrite("a.png",txt_img )
    print(txt_img.shape)
    # image_array = np.frombuffer(txt_img.tobytes(), dtype=np.uint8)
    # image_array = image_array.reshape((txt_img.size[1], txt_img.size[0], 4))
    # image_array = image_array[:,:,:3]

    gray_image = cv2.cvtColor(txt_img, cv2.COLOR_BGR2GRAY)
    bin_image = cv2.threshold(gray_image, 10, 255, cv2.THRESH_BINARY_INV)[1]
    # cv2.imwrite("a.png",bin_image )

    col_sum = np.sum(bin_image, axis=1)
    col_sum[col_sum>=1] = 1
    col_sum[col_sum<1] = 0
    col_sum = np.array(col_sum, dtype=np.int32)
    print("col_sum: ", col_sum)

    count = 0
    start_end_list = []
    start = 0
    end = 0
    for i in range(1, len(col_sum)):
        if col_sum[i] - col_sum[i-1]> 0:
            start = i
        if col_sum[i] - col_sum[i-1]< 0:
            end = i
            start_end_list.append([start, end])
    print("start_end_list: ", start_end_list)

    for s,e in start_end_list:
        crop_img = bin_image[s:e,:]
        row_sum = np.sum(crop_img, axis=0)
        row_sum = np.array(row_sum, dtype=np.int32)
        row_sum[row_sum>=1] = 1
        row_sum[row_sum<1] = 0
        
        # row_index = np.where(row_sum>0)[0]
        # row_start = row_index[0]
        # row_end = row_index[-1]
        start_t = 0
        end_t = 0
        ########## 省份切割 包含 川 特殊情况
        continue_count  = False
        for i in range(1, len(row_sum)):
            if row_sum[i] - row_sum[i-1]> 0 and not continue_count:
                start_t = i
            if row_sum[i] - row_sum[i-1]< 0:
                if i - start_t < 2:
                    continue_count = True
                    continue
                continue_count = False
                end_t = i
                crop_save_img = txt_img[s:e, start_t:end_t]
                cv2.imwrite("./crop/{}.png".format(count), crop_save_img)
                count += 1
        ########## 

        ########### 数字字母切割
        # for i in range(1, len(row_sum)):
        #     if row_sum[i] - row_sum[i-1]> 0 :
        #         start_t = i
        #     if row_sum[i] - row_sum[i-1]< 0:
        
        #         end_t = i
        #         crop_save_img = copy.copy(txt_img[s:e, start_t:end_t])
        #         print(crop_save_img.shape)
        #         if end_t - start_t == 27:
        #             crop_save_img = cv2.copyMakeBorder(crop_save_img,  0,0, 8, 8,  cv2.BORDER_CONSTANT, value=[255,255,255])

        #         cv2.imwrite("./crop/{}.png".format(count), crop_save_img)
        #         count += 1
        ########### 数字字母切割


    # return pts


def split_plate_image(img_path, diff_height=0):
    txt_img = cv2.imread(img_path)
    # txt_img = txt_img[:, 20:-20,:]

    # print(txt_img)
    # cv2.imwrite("a.png",txt_img )
    print(txt_img.shape)
    # image_array = np.frombuffer(txt_img.tobytes(), dtype=np.uint8)
    # image_array = image_array.reshape((txt_img.size[1], txt_img.size[0], 4))
    # image_array = image_array[:,:,:3]

    gray_image = cv2.cvtColor(txt_img, cv2.COLOR_BGR2GRAY)
    bin_image = cv2.threshold(gray_image, 10, 255, cv2.THRESH_BINARY_INV)[1]
    # cv2.imwrite("a.png",bin_image )

    # col_sum = np.sum(bin_image, axis=1)
    # gray_image[gray_image==0] = 0
    # gray_image[gray_image>0] = 255

    # bin_image = np.array(gray_image, dtype=np.int32)
    # print("col_sum: ", col_sum)

    col_sum = np.sum(bin_image, axis=1)
    col_sum[col_sum>=1] = 1
    col_sum[col_sum<1] = 0

    count = 0
    start_end_list = []
    start = 0
    end = 0
    for i in range(1, len(col_sum)):
        if col_sum[i] - col_sum[i-1]> 0:
            start = i
        if col_sum[i] - col_sum[i-1]< 0:
            end = i
            start_end_list.append([start, end])
    print("start_end_list: ", start_end_list)

    for s,e in start_end_list:
        crop_img = bin_image[s:e,:]
        row_sum = np.sum(crop_img, axis=0)
        row_sum = np.array(row_sum, dtype=np.int32)
        row_sum[row_sum>=1] = 1
        row_sum[row_sum<1] = 0
        
        # row_index = np.where(row_sum>0)[0]
        # row_start = row_index[0]
        # row_end = row_index[-1]
        start_t = 0
        end_t = 0

        ########### 数字字母切割
        for i in range(1, len(row_sum)):
            if row_sum[i] - row_sum[i-1]> 0 :
                start_t = i
            if row_sum[i] - row_sum[i-1]< 0:
        
                end_t = i
                crop_save_img = copy.copy(txt_img[s:e, start_t:end_t])
                print(crop_save_img.shape)
                
                cv2.imwrite("./crop/{}.png".format(count), crop_save_img)
                count += 1
        ########### 数字字母切割


    # return pts

if __name__ == "__main__":
    # img_path = r"C:\Users\swls\Desktop\1x\2.png"
    # split_image(img_path)

    img_path = r"C:\Users\swls\Desktop\1x\plate\2.png"
    split_image(img_path)